from controller import Robot
import math
import numpy as np
import copy
import matplotlib.pyplot as plt
import csv
import ast


# create the Robot instance.
robot = Robot()

# get the time step of the current world.
timestep = int(robot.getBasicTimeStep())

robot_motor = []  # LF RF RB LB
robot_motor.append(robot.getDevice("LFJ0Motor"))
robot_motor.append(robot.getDevice("LFJ1Motor"))
robot_motor.append(robot.getDevice("LFJ2Motor"))


robot_pos_sensor = []  # LF RF RB LB
# 左前足位置传感器
robot_pos_sensor.append(robot.getDevice("LFJ0Sensor"))
robot_pos_sensor.append(robot.getDevice("LFJ1Sensor"))
robot_pos_sensor.append(robot.getDevice("LFJ2Sensor"))
# 右前足位置传感器


# 触地传感器
# LF RF RB LB
robot_touch_sensor = []
robot_touch_sensor.append(robot.getDevice("LF_Touch"))





# 已测试
def webot_device_init():

    for leg in range(4):
        # 跨关节锁定
        robot_motor[0].setPosition(0)
        # 初始化使能接触传感器
        robot_touch_sensor[0].enable(timestep)
        # 初始化位置传感器
        for joint in range(3):
            robot_pos_sensor[joint].enable(timestep)


# 已测试
def set_motor_torque(name, torque):
    max = 5
    if torque > max:
        # print("error torque >max=", torque)
        torque = max
    if torque < -max:
        # print("error torque >max=", torque)
        torque = -max
    robot_motor[name].setTorque(torque)

# 顺序是 LF RF RB LB，单位是角度
def get_motor_angle(name):
    angle = 0
    angle = robot_pos_sensor[name].getValue()
    return angle * 180.0 / math.pi

def get_gyro():
    data = Gyro.getValues()
    return data

# 已测试
def get_all_motor_angle():
    temp_list = list()
    for joint in range(3):
        temp_list.append(get_motor_angle(joint))
    return temp_list

# 已测试
def is_foot_touching(leg):
    return robot_touch_sensor.getValue()

class Quadruped_robot_mix():
    def __init__(self):
        self.fall_flag = False
        self.first_run_flag = True

        self.zh = 0.1  # 抬腿到高度 0.3
        self.H_des = -0.3  # 0.5
        self.Tf = 0.3  # 飞行时间
        self.Ts = 0.3  # 支撑时间
        # 机器人机身参数
        self.l = 0.2 + 0.043 * 2  # 前后髋关节w
        self.w = 0.105  # 左右髋关节
        self.a0 = 0.0754  # 髋部连杆长度
        self.a1 = 0.2  # 大腿连杆长度
        self.a2 = 0.2  # 小腿连杆长度
        self.M = 2  # + 8 * 0.25  # 躯干质量

        ###############需要调节的参数一览#####################
        # 使用GPS的支撑相系数

        self.ST_Kz = 100
        self.ST_Bz = 1

        # 摆动相VMC参数
        self.SW_Kx = 1
        self.SW_Dx = 1
        self.SW_Ky = 1
        self.SW_Dy = 1
        self.SW_Kz = 1
        self.SW_Dz = 1
        self.Kvx = 2
        self.Kvy = 0.005

        ###############需要调节的参数一览#####################
        # 期望值
        self.Vx_des = 0
        self.Vy_des = 0
        self.Wz_des = 0
        self.Yaw_des = 0
        # 实际根据GPS得到的机身前向速度，侧向速度和yaw角转速,需要旋转矩阵
        self.Vx = 0
        self.Vy = 0
        self.Wz = 0
        self.h = 0
        # 存储变量
        # 当前欧拉角
        self.Pitch = 0
        self.Roll = 0
        self.Yaw = 0

        # 上一时刻
        self.pre_eulerAngle = [0, 0, 0]
        self.pre_pi_angle = 0
        self.pre_h = -0.32  # 初始的悬空高度，需要修改

        self.pre_vx = 0
        self.pre_vy = 0
        self.pre_wz = 0

        # 欧拉角导数
        self.dot_Pitch = 0
        self.dot_Roll = 0
        self.dot_Yaw = 0
        # 上一时刻欧拉角导数
        self.dot_Pitch_pre = 0
        self.dot_Roll_pre = 0
        self.dot_Yaw_pre = 0
        # 四脚着地状态 逆时针编号
        self.is_touching = 0
        # 四足足端当前坐标
        self.pos = [0, 0, 0]
        # 四足关节角，有12个
        self.theta = list()
        # 四足足端当前导数
        self.dot_pos = [0,0,0]
        # 四足足端上一时坐标
        self.pre_pos = [0, 0, 0]
        # 步态时间
        self.time = 0.0

    # 机身腿部坐标转换到单腿坐标下，是机身坐标，不是世界系坐标
    def convert_pos(self, leg, pos):
        x, y, z = pos
        if leg == 0 or leg == 1:
            x -= self.l / 2
        else:
            x += self.l / 2

        # if leg == 0 or leg == 3:
        #     y -= self.w / 2
        # else:
        #     y += self.w / 2
        return [x, y, z]

    # 这里出来的坐标是机身坐标系，搞清楚电机转的顺序，这里都是顺时针为正转，腿的顺序，腿的顺序LF RF RB LB
    def forwardkinematics(self, theta):
        a0 = self.a0  # 腿长
        a1 = - self.a1  # 腿长
        a2 = - self.a2  # 腿长

        flag = [1.0, -1.0, -1.0, 1.0]
        # 这个和腿的顺序有关
        flag_d = [1.0, 1.0, -1.0, -1.0]
        # 这个和腿的顺序有关
        flag_l = [1.0, -1.0, -1.0, 1.0]
        # 足底位置

        temp_toe = list()

        s0 = math.sin(theta[0] / 180.0 * math.pi)
        c0 = math.cos(theta[0] / 180.0 * math.pi)
        s1 = math.sin(theta[1] / 180.0 * math.pi)
        c1 = math.cos(theta[1] / 180.0 * math.pi)

        s12 = math.sin(theta[1] / 180.0 * math.pi + theta[2] / 180.0 * math.pi)
        c12 = math.cos(theta[1] / 180.0 * math.pi + theta[2] / 180.0 * math.pi)

        x = a1 * s1 + a2 * s12
        y = a0 * c0 - a1 * c1 * s0 - a2 * c12 * s0
        z = a0 * s0 + a1 * c1 * c0 + a2 * c12 * c0

        temp_toe.append(x)
        temp_toe.append(y)
        temp_toe.append(z)

        return temp_toe

    # 只传入一个腿的三个关节值
    def create_transJ(self, theta, leg):
        a0 = self.a0
        a1 = - self.a1  # 腿长
        a2 = - self.a2  # 腿长

        flag = [1.0, -1.0, -1.0, 1.0]

        s0 = math.sin(theta[0] / 180.0 * math.pi)
        c0 = math.cos(theta[0] / 180.0 * math.pi)
        s1 = math.sin(theta[1] / 180.0 * math.pi)
        c1 = math.cos(theta[1] / 180.0 * math.pi)

        s12 = math.sin(theta[1] / 180.0 * math.pi + theta[2] / 180.0 * math.pi)
        c12 = math.cos(theta[1] / 180.0 * math.pi + theta[2] / 180.0 * math.pi)

        transJ = np.zeros((3, 3))
        # 分别求偏导数
        transJ[0, 0] = 0
        transJ[0, 1] = a1 * c1 + a2 * c12
        transJ[0, 2] = a2 * c12

        transJ[1, 0] = -flag[leg] * a0 * s0 - (a1 * c1 + a2 * c12) * c0
        transJ[1, 1] = s0 * (a1 * s1 + a2 * s12)
        transJ[1, 2] = a2 * s0 * s12

        transJ[2, 0] = - s0 * (a1 * c1 + a2 * c12) + flag[leg] * a0 * c0
        transJ[2, 1] = - c0 * (a1 * s1 + a2 * s12)
        transJ[2, 2] = - a2 * c0 * s12
        return transJ.T

    def update_theta(self):
        self.theta = get_all_motor_angle()



    # 更新全部状态
    def update_robot_state(self):

        # 更新关节角
        self.update_theta()
        # 更新四条腿的运动学正解
        self.pos = self.forwardkinematics(self.theta)
        # 跟新足底速度
        self.time += timestep * 0.001



    def control(self):
        # 支持相控制----------------------------------------
        h_des = self.H_des
        self.h = self.pos[2]
        dh = (self.h - self.pre_h) / (0.001 * timestep)
        self.pre_h = self.h

        F = np.zeros((3, 1))

        # 求支撑力，两条腿，符号不一样
        Gx_ff = self.M * 9.8 * math.sin(self.Pitch / 180.0 * math.pi)
        Gy_ff = self.M * 9.8 * math.sin(self.Roll / 180.0 * math.pi)
        Gz_ff = self.M * 9.8 * math.cos(self.Pitch / 180.0 * math.pi) * math.cos(self.Roll / 180.0 * math.pi)


        F[0][0] = 0

        F[1][0] = 0

        F[2][0] = self.ST_Kz * (h_des - self.h) + self.ST_Bz * (0 - dh)  # + Gz_ff / 2


        tau_left = self.create_transJ(self.theta, 0)
        tau = - tau_left @ F

        print(tau)
        set_motor_torque(0, tau[0][0])
        set_motor_torque(1, tau[1][0])
        set_motor_torque(2, tau[2][0])

    def run(self):
        # 初始化webots节点设备
        webot_device_init()
        while robot.step(timestep) != -1:
            # 更新数据
            self.update_robot_state()
            self.control()


if __name__ == '__main__':
    q = Quadruped_robot_mix()
    q.run()
